/* Copyright 2025 The OpenXLA Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#ifndef XLA_TESTS_XLA_TEST_BACKEND_PREDICATES_H_
#define XLA_TESTS_XLA_TEST_BACKEND_PREDICATES_H_

#include <vector>

#include <gtest/gtest.h>
#include "absl/base/nullability.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"

namespace xla::test {
// The utilities in this file are used to disable test cases for targets in
// targets generated by `xla_test`. Example usage:
// ```
// TEST(MyXlaTest, MyXlaTestcase) {
//    if (xla::test::DeviceIs(xla::test::kA100)) {
//      GTEST_SKIP();
//    }
//    ...
// }
//
// TEST(MyXlaTest, MyXlaTestcase) {
//    if (xla::test::DeviceTypeIs(xla::test::kCpu) {
//      GTEST_SKIP();
//    }
//    ...
// }
// ```
// For a given backend, the device and modifiers are given like
// `device, *modifiers = "_".split(backend)`
// Note that for GPU this means that `a100` or similar would be considered a
// modifier rather than a device name.
//
// To see how xla_test sets the device and modifiers concretely, use
// `bazel query $TARGET_NAME --output=build`.

// Constants to avoid typos :-).
inline constexpr const absl::string_view kCpu = "cpu";
inline constexpr const absl::string_view kGpu = "gpu";
inline constexpr const absl::string_view kA100 = "a100";
inline constexpr const absl::string_view kH100 = "h100";
inline constexpr const absl::string_view kB100 = "b100";
inline constexpr const absl::string_view kP100 = "p100";
inline constexpr const absl::string_view kV100 = "v100";

inline constexpr const absl::string_view kInterpreter = "interpreter";

inline constexpr const absl::string_view kTpu = "tpu";
inline constexpr const absl::string_view kHardware = "hardware";
inline constexpr const absl::string_view kIss = "iss";
inline constexpr const absl::string_view kGrm = "grm";

absl::string_view GetEnvOrDie(const char* absl_nonnull key);

// Returns true if `device` matches the device under test.
bool DeviceIs(absl::string_view device);

bool DeviceIsOneOf(absl::Span<const absl::string_view> devices);

// Matches based on wider `device_type`, like `tpu|gpu|cpu|interpreter`.
bool DeviceTypeIs(absl::string_view device_type);

bool DeviceTypeIsOneOf(absl::Span<const absl::string_view> devices);

// Returns true if all `modifiers` match the current backend.
bool HasModifiers(absl::Span<const absl::string_view> modifiers);

// Returns true if `device` matches the device under test and all modifiers
// match the current backend. Like `DeviceIs(device) && HasModifiers(modifiers)`
bool BackendLike(absl::string_view device,
                 absl::Span<const absl::string_view> modifiers);

// Same as `BackendLike`, but modifiers must be an exact match.
bool BackendIsExactly(absl::string_view device,
                      absl::Span<const absl::string_view> modifiers);

// Returns true only for base variant hardware + emulation.
bool BackendIsStrict(absl::string_view device);

bool BackendSupportsFloat64();
bool BackendSupportsComplex128();

bool UsingStreamExecutorGpuClient();

// Useful to generate an intentionally empty set of inputs for a parameterized
// test. This is needed when we are manipulating the inputs based on the
// backend and would like some backends to receive zero inputs. Usage of this
// function should always be accompanied by
// `GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST()`.
//
// Example:
//
// GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(TestFixture);
// INSTANTIATE_TEST_SUITE_P(
//    ParameterizedTestName,
//    TestFixture,
//    DeviceIs(kA100) ? ::testing::ValuesIn({1, 2, 3}) : Empty<int>());
template <typename T>
::testing::internal::ParamGenerator<T> Empty() {
  std::vector<T> empty_vec;
  return ::testing::ValuesIn(empty_vec);
}

}  // namespace xla::test
#endif  // XLA_TESTS_XLA_TEST_BACKEND_PREDICATES_H_
